#include <iostream>
#include <thread>
#include <vector>
#include <queue>
#include <mutex>
#include <condition_variable>
#include <chrono>
#include <optional>

struct Task {
    int id;
    int payload;
};

// 一个线程安全的任务队列，支持 Master 推送任务和 Slave 拉取任务
class TaskQueue {
public:
    void push(Task task) {
        {
            std::lock_guard<std::mutex> lock(mutex_);
            tasks_.push(std::move(task));
        }
        cv_.notify_one();
    }

    // 尝试获取任务，如果接收到停止信号则返回 std::nullopt
    std::optional<Task> pop() {
        std::unique_lock<std::mutex> lock(mutex_);
        cv_.wait(lock, [&] { return stop_ || !tasks_.empty(); });

        if (stop_ && tasks_.empty()) {
            return std::nullopt;
        }

        Task task = std::move(tasks_.front());
        tasks_.pop();
        return task;
    }

    // 发送停止信号，通知所有等待的线程退出
    void stop() {
        {
            std::lock_guard<std::mutex> lock(mutex_);
            stop_ = true;
        }
        cv_.notify_all();
    }

private:
    std::queue<Task> tasks_;
    std::mutex mutex_;
    std::condition_variable cv_;
    bool stop_{false};
};

// 模拟耗时计算
int process(int input) {
    std::this_thread::sleep_for(std::chrono::milliseconds(100));
    return input * input;
}

int main() {
    TaskQueue task_queue;
    const int num_workers = 4;

    // 启动 Slave 工作线程
    std::vector<std::thread> workers;
    for (int i = 0; i < num_workers; ++i) {
        workers.emplace_back([i, &task_queue] {
            while (true) {
                auto task_opt = task_queue.pop();
                if (!task_opt.has_value()) {
                    std::cout << "[Worker " << i << "] 收到退出信号。\n";
                    break;
                }

                Task task = task_opt.value();
                int result = process(task.payload);
                std::cout << "[Worker " << i << "] 处理任务 " << task.id
                          << "，输入 " << task.payload
                          << "，输出 " << result << '\n';
            }
        });
    }

    // Master 线程：生成任务
    const int total_tasks = 10;
    for (int i = 0; i < total_tasks; ++i) {
        Task task{i, i + 1};
        std::cout << "[Master] 发布任务 " << task.id
                  << "，payload = " << task.payload << '\n';
        task_queue.push(task);
        std::this_thread::sleep_for(std::chrono::milliseconds(50));
    }

    // 所有任务发布完毕后，通知 workers 停止
    task_queue.stop();

    // 等待所有 worker 退出
    for (auto& worker : workers) {
        worker.join();
    }

    std::cout << "[Master] 所有任务处理完毕，程序结束。\n";
    return 0;
}